from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model
import torch
from datasets import load_from_disk
from transformers import DataCollatorForLanguageModeling

# 参数配置
model_path = "./models/deepseek-ai/DeepSeek-R1-Distill-Qwen-1___5B"
tokenizer = AutoTokenizer.from_pretrained(model_path)
tokenizer.pad_token = tokenizer.eos_token  # 确保pad_token已设置

# 加载数据
dataset = load_from_disk("./data/processed_data")

print("训练集样本数:", len(dataset["train"]))  # 应为3
print("验证集样本数:", len(dataset["validation"]))  # 应为3
print("首条训练数据样例:", dataset["train"][0])  # 查看具体字段

# PEFT配置
peft_config = LoraConfig(
    r=4,
    lora_alpha=8,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.3,
    bias="none",
    task_type="CAUSAL_LM"
)

# 训练参数（修正参数名）
training_args = TrainingArguments(
    output_dir="./output",
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=8,
    learning_rate=2e-5,
    num_train_epochs=3,
    logging_steps=10,
    eval_strategy="epoch",  # 替换为eval_strategy
    save_strategy="epoch",
    # gradient_checkpointing=True,  # 启用梯度检查点（节省显存）
    fp16=True,
    report_to=["tensorboard"]
)

# 初始化模型
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    attn_implementation="eager",  # 关闭sdpa优化
    torch_dtype=torch.float16,
    device_map="auto"
)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

# 测试显存是否足够加载模型
# dummy_input = torch.randint(0, 10000, (1, 16)).to("cuda")  # 模拟16个token的输入
# with torch.no_grad():
#     outputs = model(dummy_input)  # 若此处报错，说明显存不足
# print("显存测试通过")

# 创建数据整理器
data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,  # 因果语言模型使用mlm=False
)

# 配置Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["validation"],
    data_collator=data_collator,  # 使用新整理器
)

# 开始训练
trainer.train()

# 保存模型
model.save_pretrained("./finetuned_model")
